"""
Message system for agent communication
"""
from dataclasses import dataclass
from enum import Enum, auto
from typing import Any, Optional, Dict


class MessageType(Enum):
    """消息类型枚举"""
    TEXT = auto()  # 普通文本消息
    TOOL_CALL = auto()  # 工具调用请求
    TOOL_RESULT = auto()  # 工具调用结果
    CONTEXT = auto()  # 上下文相关消息


@dataclass
class Message:
    """消息基类"""
    content: Any
    sender: str
    type: MessageType = MessageType.TEXT
    receiver: Optional[str] = None  # None表示广播消息
    metadata: Optional[Dict] = None

    def to_dict(self) -> Dict:
        """将消息转换为字典格式"""
        return {
            "type": self.type.name,
            "content": self.content,
            "sender": self.sender,
            "receiver": self.receiver,
            "metadata": self.metadata or {}
        }

    @classmethod
    def from_dict(cls, data: Dict) -> 'Message':
        """从字典创建消息实例"""
        return cls(type=MessageType[data["type"]],
                   content=data["content"],
                   sender=data["sender"],
                   receiver=data.get("receiver"),
                   metadata=data.get("metadata", {}))


@dataclass
class ToolCallMessage(Message):
    """工具调用消息"""

    def __init__(self,
                 tool_name: str,
                 arguments: Dict,
                 sender: str,
                 receiver: Optional[str] = None,
                 metadata: Optional[Dict] = None):
        super().__init__(type=MessageType.TOOL_CALL,
                         content={
                             "tool_name": tool_name,
                             "arguments": arguments
                         },
                         sender=sender,
                         receiver=receiver,
                         metadata=metadata)


@dataclass
class ToolResultMessage(Message):
    """工具调用结果消息"""

    def __init__(self,
                 success: bool,
                 result: Any,
                 error: Optional[str] = None,
                 sender: str = "system",
                 receiver: Optional[str] = None,
                 metadata: Optional[Dict] = None):
        super().__init__(type=MessageType.TOOL_RESULT,
                         content={
                             "success": success,
                             "result": result,
                             "error": error
                         },
                         sender=sender,
                         receiver=receiver,
                         metadata=metadata)
